import socket
import json
import time
import zmq
import threading
from typing import Dict
from common import *
from multiprocessing import Manager

class ThreadDeviceServer:
    def __init__(self, shared_devices: Dict[str, dict] = None):
        self.manager = Manager() if shared_devices is None else None
        self.devices: Dict[str, dict] = shared_devices if shared_devices is not None else self.manager.dict()
        self.context = zmq.Context()
        self.sub_socket = self.context.socket(zmq.SUB)
        self.sub_socket.setsockopt_string(zmq.SUBSCRIBE, "")
        self.connected = set()
        self.running = True
        self.threads = []
        self._lock = threading.Lock()

    def udp_listener(self):
        udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udp_socket.bind(("", BROADCAST_PORT))

        while self.running:
            try:
                data, addr = udp_socket.recvfrom(1024)
                info = json.loads(data.decode())
                uuid_ = info['uuid']
                info['ip'] = addr[0]
                info['last_seen'] = time.time()
                info['status'] = 'online'

                with self._lock:
                    if uuid_ not in self.devices:
                        self.devices[uuid_] = info
                        print(f"[UDP] 发现设备 {info['name']} @ {addr[0]}")
                    else:
                        self.devices[uuid_].update(info)
            except Exception as e:
                if self.running:  # 只在正常运行时打印错误
                    print("[UDP] 数据解析失败：", e)

    def zmq_subscriber(self):
        retry_counts = {}
        while self.running:
            with self._lock:
                devices_copy = dict(self.devices)

            for uuid_, info in devices_copy.items():
                addr = f"tcp://{info['ip']}:{info['pub_port']}"
                if addr not in self.connected:
                    try:
                        if uuid_ not in retry_counts:
                            retry_counts[uuid_] = 0
                        if retry_counts[uuid_] < MAX_RETRY_COUNT:
                            self.sub_socket.connect(addr)
                            self.connected.add(addr)
                            print(f"[ZMQ] 订阅设备 {info['name']} 的心跳")
                            retry_counts[uuid_] = 0
                        else:
                            print(f"[警告] 设备 {info['name']} 连接失败次数过多，等待下次发现")
                            retry_counts[uuid_] = 0
                            del self.devices[uuid_]
                            self.sub_socket.disconnect(addr)
                            self.connected.remove(addr)

                    except zmq.ZMQError as e:
                        retry_counts[uuid_] += 1
                        print(f"[ZMQ] 连接设备 {info['name']} 失败: {e}")
                        time.sleep(RECONNECT_DELAY)

            try:
                while self.running:
                    try:
                        msg = self.sub_socket.recv_string(flags=zmq.NOBLOCK)
                        heartbeat = json.loads(msg)
                        uuid_ = heartbeat['uuid']
                        if uuid_ in self.devices:
                            with self._lock:
                                device_info = self.devices[uuid_]
                                device_info['last_seen'] = time.time()
                                if self.devices[uuid_]['status'] != 'online':
                                    print(f"[恢复] 设备 {self.devices[uuid_]['name']} 恢复在线")
                                device_info['status'] = 'online'
                                self.devices[uuid_] = device_info
                            retry_counts[uuid_] = 0
                    except zmq.Again:
                        break
            except Exception as e:
                if self.running:
                    print(f"[ZMQ] 接收心跳异常: {e}")
            time.sleep(HEARTBEAT_INTERVAL)

    def device_status_checker(self):
        status_change_time = {}
        offline_count = {}
        while self.running:
            try:
                now = time.time()
                with self._lock:
                    devices_copy = dict(self.devices)

                for uuid_, info in devices_copy.items():
                    time_diff = now - info.get('last_seen', 0)
                    current_status = info['status']
                    
                    new_status = current_status
                    if time_diff > OFFLINE_TIMEOUT:
                        if uuid_ not in offline_count:
                            offline_count[uuid_] = 0
                        offline_count[uuid_] += 1
                        
                        if offline_count[uuid_] >= 3:
                            new_status = 'offline'
                    else:
                        offline_count[uuid_] = 0
                        new_status = 'online'
                    
                    if new_status != current_status:
                        last_change = status_change_time.get(uuid_, 0)
                        if now - last_change > RECONNECT_DELAY:
                            with self._lock:
                                self.devices[uuid_]['status'] = new_status
                            status_change_time[uuid_] = now
                            
                            if new_status == 'offline':
                                print(f"[警告] 设备 {info['name']} 掉线")
                                addr = f"tcp://{info['ip']}:{info['pub_port']}"
                                if addr in self.connected:
                                    try:
                                        self.sub_socket.disconnect(addr)
                                        self.connected.remove(addr)
                                        print(f"[ZMQ] 断开与设备 {info['name']} 的连接")
                                    except zmq.ZMQError as e:
                                        print(f"[ZMQ] 断开连接失败: {e}")
                            else:
                                print(f"[恢复] 设备 {info['name']} 重新上线")
            except Exception as e:
                if self.running:
                    print(f"[错误] 设备状态检查异常: {e}")
            time.sleep(HEARTBEAT_INTERVAL)

    def get_online_devices(self):
        now = time.time()
        with self._lock:
            return [
                {
                    "uuid": d["uuid"],
                    "name": d["name"],
                    "ip": d["ip"],
                    "info": d.get("info", "")
                }
                for d in self.devices.values()
                if now - d.get("last_seen", 0) < 10
            ]

    def start(self):
        print("[Server] 启动设备监听服务")
        try:
            # 创建并启动线程
            udp_thread = threading.Thread(target=self.udp_listener)
            zmq_thread = threading.Thread(target=self.zmq_subscriber)
            status_thread = threading.Thread(target=self.device_status_checker)
            
            self.threads = [udp_thread, zmq_thread, status_thread]
            
            for thread in self.threads:
                thread.daemon = True
                thread.start()
            
            # 等待线程结束
            for thread in self.threads:
                thread.join()
                
        except Exception as e:
            print(f"[错误] 服务器运行异常: {e}")

    def stop(self):
        self.running = False
        # 等待所有线程结束
        for thread in self.threads:
            thread.join(timeout=5)

if __name__ == "__main__":
    server = ThreadDeviceServer()
    try:
        server.start()
    except KeyboardInterrupt:
        server.stop()